{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ResNet和Transformer的Lipschitz常量"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 定义残差网络 ResNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "import numpy as np\n",
    "import random\n",
    "from os import name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置随机种子\n",
    "torch.manual_seed(0)\n",
    "random.seed(0)\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置残差结构\n",
    "class BasicBlock(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, stride=1):\n",
    "        super(BasicBlock, self).__init__()\n",
    "\n",
    "        # 卷积层1\n",
    "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)\n",
    "        # batch normalize 1\n",
    "        self.bn1 = nn.BatchNorm2d(out_channels)\n",
    "        # 卷积层2\n",
    "        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)\n",
    "        # batch normalize 2\n",
    "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
    "\n",
    "        # 连接边\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_channels != out_channels:\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(out_channels)\n",
    "            )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = torch.relu(out)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "\n",
    "        out += self.shortcut(x)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置resnet网络\n",
    "class ResNet(nn.Module):\n",
    "    def __init__(self, block, num_blocks, dim):\n",
    "        super(ResNet, self).__init__()\n",
    "\n",
    "        self.num_blocks = num_blocks # 残差模块重复的次数\n",
    "        self.in_channels = dim       # 输入维度\n",
    "\n",
    "        self.blocks = []\n",
    "        for i in range(num_blocks):\n",
    "            self.blocks.append(BasicBlock(dim, dim))\n",
    "\n",
    "    def forward(self, x):\n",
    "        tmp_x = x\n",
    "        for i in range(self.num_blocks ):\n",
    "            tmp_x = self.blocks[i](tmp_x)\n",
    "        return tmp_x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### transformer中的Self-Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, embed_dim, num_heads):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_heads = num_heads\n",
    "        \n",
    "        assert embed_dim % num_heads == 0\n",
    "        self.head_dim = embed_dim // num_heads\n",
    "\n",
    "        self.q = nn.Linear(embed_dim, embed_dim)\n",
    "        self.k = nn.Linear(embed_dim, embed_dim)\n",
    "        self.v = nn.Linear(embed_dim, embed_dim)\n",
    "\n",
    "        self.fc_out = nn.Linear(embed_dim, embed_dim)\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        batch_size = x.shape[0]\n",
    "\n",
    "        q = self.q(x).view(batch_size, -1, self.num_heads, self.head_dim)\n",
    "        k = self.k(x).view(batch_size, -1, self.num_heads, self.head_dim)\n",
    "        v = self.v(x).view(batch_size, -1, self.num_heads, self.head_dim)\n",
    "        q = q.permute(0,2,1,3)\n",
    "        k = k.permute(0,2,1,3)\n",
    "        v = v.permute(0,2,1,3)\n",
    "\n",
    "        query,key,value = q,k,v\n",
    "        energy = torch.einsum(\"nqhd,nkhd->nhqk\", [query,key])\n",
    "\n",
    "        if mask is not None:\n",
    "            energy = energy.masked_fill(mask == 0, float(\"-1e20\"))\n",
    "\n",
    "        attention = torch.softmax(energy / (self.embed_dim ** (1/2)), dim=-1)\n",
    "        out = torch.einsum(\"nhql,nlhd->nqhd\", [attention, value]).permute(0,2,1,3).contiguous()\n",
    "        out = out.view(batch_size, -1, self.embed_dim)\n",
    "\n",
    "        return self.fc_out(out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, embed_dim, num_heads, forward_expansion=4):\n",
    "        super(TransformerBlock, self).__init__()\n",
    "\n",
    "        self.attention = MultiHeadAttention(embed_dim, num_heads)\n",
    "        self.norm1 = nn.LayerNorm(embed_dim)\n",
    "        self.norm2 = nn.LayerNorm(embed_dim)\n",
    "\n",
    "        self.feed_forward = nn.Sequential(\n",
    "            nn.Linear(embed_dim, forward_expansion * embed_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(forward_expansion * embed_dim, embed_dim)\n",
    "        )\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        attention = self.attention(x, mask)\n",
    "        x = self.norm1(attention + x)\n",
    "\n",
    "        forward = self.feed_forward(x)\n",
    "        out = self.norm2(forward + x)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, embed_dim, num_heads, num_layers, forward_expansion):\n",
    "        super(Transformer, self).__init__()\n",
    "\n",
    "        self.layers = nn.ModuleList(\n",
    "            [TransformerBlock(embed_dim, num_heads,forward_expansion) for _ in range(num_layers)]\n",
    "        )\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        out = x\n",
    "        for layer in self.layers:\n",
    "            out = layer(out, mask)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import transformers.models.bert"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 仿真实验\n",
    "结果：\n",
    "1. 残差网络：随着网络层数的增加，Lipschitz常量变化还可控，但transformer会增长的特别快\n",
    "2. Lipschitz常量越大，表达能力越强\n",
    "3. 参数的调整对transformer的影响比较大\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ResNet_L(L, dim):\n",
    "    return ResNet(BasicBlock, L, dim)\n",
    "\n",
    "def init_weights(m):\n",
    "    '''\n",
    "    初始化权重\n",
    "    '''\n",
    "    # 判断是否为线性层或者卷积层\n",
    "    if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):\n",
    "        torch.nn.init.xavier_uniform_(m.weight, gain=factor)\n",
    "        m.bias.data.fill_(0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.对比resnet和transformer网络的Lipschitz常量【未训练】"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of Layers:  1     resnet, transformer:  tensor(10.0514) tensor(5.6411)\n",
      "Number of Layers:  2     resnet, transformer:  tensor(21.1816) tensor(10.8477)\n",
      "Number of Layers:  4     resnet, transformer:  tensor(42.5665) tensor(30.6373)\n",
      "Number of Layers:  8     resnet, transformer:  tensor(81.4569) tensor(188.6219)\n",
      "Number of Layers:  12     resnet, transformer:  tensor(116.2124) tensor(1167.3009)\n",
      "Number of Layers:  16     resnet, transformer:  tensor(149.4668) tensor(6967.8198)\n"
     ]
    }
   ],
   "source": [
    "# 使用例子\n",
    "# 参数设置\n",
    "embed_dim = 512   # 通道数\n",
    "num_heads = 8\n",
    "forward_factor = 4\n",
    "\n",
    "K = 32  # 图片大小\n",
    "B = 10  # batch size\n",
    "factor = 2.0  # xavier初始化的gain值\n",
    "episilon_factor = 1e-7    # 随机噪声\n",
    "\n",
    "# 存放网络结果\n",
    "result_12_resnet = []\n",
    "result_12_dp = []\n",
    "\n",
    "# 仿真输入\n",
    "x = torch.randn(B, embed_dim, K, K)    # 图片\n",
    "episilon = episilon_factor * torch.randn(B, embed_dim, K, K)   # 噪声\n",
    "\n",
    "#layers = [1,2,4,8,12,16,24,32,48,64]   # 层次\n",
    "layers = [1,2,4,8,12,16] \n",
    "for L in layers:\n",
    "    # 分别定义2个网络，并进行初始化\n",
    "    transformer = Transformer(embed_dim, num_heads=8, num_layers=L, forward_expansion=4)\n",
    "    transformer.apply(init_weights)     #--【应用函数】\n",
    "\n",
    "    resnet = ResNet(BasicBlock, L, embed_dim)\n",
    "    resnet.apply(init_weights)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        # x: 10*512*32*32  #批量，通道，大小【resnet输入】\n",
    "        # x: 10*1024*512   #transformer输入\n",
    "\n",
    "        # 1.输入x1和x2得到y1和y2\n",
    "        # x1: 原图\n",
    "        x1 = x\n",
    "        y1_resnet = resnet(x1)              # resnet结果\n",
    "\n",
    "        x1 = torch.reshape(x, [B, embed_dim, K*K])\n",
    "        x1 = x1.transpose(2, 1)             # 输入：10, 1024, 512\n",
    "        y1_transformer = transformer(x1)    # transformer结果\n",
    "\n",
    "        # x2: 加了扰动后的图\n",
    "        x2 = x+episilon\n",
    "        y2_resnet = resnet(x2)\n",
    "\n",
    "        x2 = torch.reshape(x2, [B, embed_dim, K*K])\n",
    "        x2 = x2.transpose(2, 1)\n",
    "        y2_transformer = transformer(x2)\n",
    "\n",
    "        # 2.计算Lipschitz常量\n",
    "        # 将参数展平【要求修改网络设计，保证输出与输入x的维度是一致的】\n",
    "        x1 = torch.reshape(x1, [B, embed_dim*K*K])\n",
    "        x2 = torch.reshape(x2, [B, embed_dim*K*K])\n",
    "        y1_resnet = torch.reshape(y1_resnet, [B, embed_dim*K*K])\n",
    "        y2_resnet = torch.reshape(y2_resnet, [B, embed_dim*K*K])\n",
    "        y1_transformer = torch.reshape(y1_transformer, [B, embed_dim*K*K])\n",
    "        y2_transformer = torch.reshape(y2_transformer, [B, embed_dim*K*K])\n",
    "\n",
    "        # 计算梯度的最大值, 即k0\n",
    "        lips_l2_res = torch.max(torch.sqrt(torch.sum(torch.pow(y2_resnet-y1_resnet, 2), axis = 1) / \n",
    "                                           torch.sum(torch.pow(x2-x1, 2), axis = 1)))\n",
    "        lips_l2_transformer = torch.max(torch.sqrt(torch.sum(torch.pow(y2_transformer-y1_transformer, 2), axis = 1) / \n",
    "                                                   torch.sum(torch.pow(x2-x1, 2), axis = 1)))\n",
    "        \n",
    "        # 将不同网络层的K0存放到结果列表中\n",
    "        result_12_resnet.append(lips_l2_res.data.numpy())\n",
    "        result_12_dp.append(lips_l2_transformer.data.numpy())\n",
    "\n",
    "        print('Number of Layers: ', L, '    resnet, transformer: ', lips_l2_res.data, lips_l2_transformer.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10.051371,21.181585,42.56654,81.45695,116.21245,149.46684\n"
     ]
    }
   ],
   "source": [
    "# 打印列表结果\n",
    "print(*np.array(result_12_resnet), sep=',')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.641077,10.847671,30.637302,188.62192,1167.3009,6967.82\n"
     ]
    }
   ],
   "source": [
    "print(*np.array(result_12_dp), sep=',')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulation_lips(embed_dim, num_heads,forward_factor,\n",
    "                    K,B,factor,episilon_factor,\n",
    "                    layers = [1,2,4,8,12,16]) -> list:\n",
    "    '''\n",
    "    调整参数进行仿真\n",
    "    '''\n",
    "    # 存放网络结果\n",
    "    result_12_resnet = []\n",
    "    result_12_dp = []\n",
    "    # 仿真输入\n",
    "    x = torch.randn(B, embed_dim, K, K)    # 图片\n",
    "    episilon = episilon_factor * torch.randn(B, embed_dim, K, K)   # 噪声\n",
    "\n",
    "    #layers = [1,2,4,8,12,16,24,32]   # 层次\n",
    "    #layers = [1,2,4]\n",
    "    for L in layers:\n",
    "            # 分别定义2个网络，并进行初始化\n",
    "        transformer = Transformer(embed_dim, num_heads=8, num_layers=L, forward_expansion=4)\n",
    "        transformer.apply(init_weights)     #--【应用函数】\n",
    "\n",
    "        resnet = ResNet(BasicBlock, L, embed_dim)\n",
    "        resnet.apply(init_weights)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            # x: 10*512*32*32  #批量，通道，大小【resnet输入】\n",
    "            # x: 10*1024*512   #transformer输入\n",
    "\n",
    "            # 1.输入x1和x2得到y1和y2\n",
    "            # x1: 原图\n",
    "            x1 = x\n",
    "            y1_resnet = resnet(x1)              # resnet结果\n",
    "\n",
    "            x1 = torch.reshape(x, [B, embed_dim, K*K])\n",
    "            x1 = x1.transpose(2, 1)             # 输入：10, 1024, 512\n",
    "            y1_transformer = transformer(x1)    # transformer结果\n",
    "\n",
    "            # x2: 加了扰动后的图\n",
    "            x2 = x+episilon\n",
    "            y2_resnet = resnet(x2)\n",
    "\n",
    "            x2 = torch.reshape(x2, [B, embed_dim, K*K])\n",
    "            x2 = x2.transpose(2, 1)\n",
    "            y2_transformer = transformer(x2)\n",
    "\n",
    "            # 2.计算Lipschitz常量\n",
    "            # 将参数展平【要求修改网络设计，保证输出与输入x的维度是一致的】\n",
    "            x1 = torch.reshape(x1, [B, embed_dim*K*K])\n",
    "            x2 = torch.reshape(x2, [B, embed_dim*K*K])\n",
    "            y1_resnet = torch.reshape(y1_resnet, [B, embed_dim*K*K])\n",
    "            y2_resnet = torch.reshape(y2_resnet, [B, embed_dim*K*K])\n",
    "            y1_transformer = torch.reshape(y1_transformer, [B, embed_dim*K*K])\n",
    "            y2_transformer = torch.reshape(y2_transformer, [B, embed_dim*K*K])\n",
    "\n",
    "            # 计算梯度的最大值, 即k0\n",
    "            lips_l2_res = torch.max(torch.sqrt(torch.sum(torch.pow(y2_resnet-y1_resnet, 2), axis = 1) / \n",
    "                                            torch.sum(torch.pow(x2-x1, 2), axis = 1)))\n",
    "            lips_l2_transformer = torch.max(torch.sqrt(torch.sum(torch.pow(y2_transformer-y1_transformer, 2), axis = 1) / \n",
    "                                                    torch.sum(torch.pow(x2-x1, 2), axis = 1)))\n",
    "            \n",
    "            # 将不同网络层的K0存放到结果列表中\n",
    "            result_12_resnet.append(lips_l2_res.data.numpy())\n",
    "            result_12_dp.append(lips_l2_transformer.data.numpy())\n",
    "\n",
    "            print('Number of Layers: ', L, '    resnet, transformer: ', lips_l2_res.data, lips_l2_transformer.data)\n",
    "    \n",
    "    return result_12_resnet, result_12_dp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. 调整xavier初始化的gain值\n",
    "2.0 -> 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of Layers:  1     resnet, transformer:  tensor(9.9579) tensor(2.3556)\n",
      "Number of Layers:  2     resnet, transformer:  tensor(21.1754) tensor(3.3262)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of Layers:  4     resnet, transformer:  tensor(42.7032) tensor(4.9326)\n",
      "Number of Layers:  8     resnet, transformer:  tensor(80.9486) tensor(7.5835)\n",
      "Number of Layers:  12     resnet, transformer:  tensor(116.1752) tensor(9.3964)\n",
      "Number of Layers:  16     resnet, transformer:  tensor(149.2384) tensor(11.0455)\n"
     ]
    }
   ],
   "source": [
    "# 使用例子\n",
    "# 参数设置\n",
    "embed_dim = 512   # 通道数\n",
    "num_heads = 8\n",
    "forward_factor = 4\n",
    "\n",
    "K = 32  # 图片大小\n",
    "B = 10  # batch size\n",
    "factor = 1.0  # ----------------------------xavier初始化的gain值\n",
    "episilon_factor = 1e-7    # 随机噪声\n",
    "\n",
    "# 存放网络结果\n",
    "result_12_resnet = []\n",
    "result_12_dp = []\n",
    "\n",
    "result_12_resnet, result_12_dp = simulation_lips(embed_dim, num_heads,forward_factor,\n",
    "                                                 K,B,factor,episilon_factor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9.957892,21.17539,42.70322,80.94862,116.17525,149.23842\n",
      "2.3556037,3.3262262,4.9325976,7.583474,9.396358,11.045459\n"
     ]
    }
   ],
   "source": [
    "# 打印列表结果\n",
    "print(*np.array(result_12_resnet), sep=',')\n",
    "print(*np.array(result_12_dp), sep=',')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.增加通道数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of Layers:  1     resnet, transformer:  tensor(12.7921) tensor(2.4456)\n",
      "Number of Layers:  2     resnet, transformer:  tensor(28.4037) tensor(3.4856)\n",
      "Number of Layers:  4     resnet, transformer:  tensor(58.2677) tensor(5.1532)\n",
      "Number of Layers:  8     resnet, transformer:  tensor(112.2736) tensor(7.6673)\n"
     ]
    }
   ],
   "source": [
    "# 使用例子\n",
    "# 参数设置\n",
    "embed_dim = 1024   # ----------------------------通道数\n",
    "num_heads = 8\n",
    "forward_factor = 4\n",
    "\n",
    "K = 32  # 图片大小\n",
    "B = 10  # batch size\n",
    "factor = 1.0  # xavier初始化的gain值\n",
    "episilon_factor = 1e-7    # 随机噪声\n",
    "\n",
    "# 存放网络结果\n",
    "result_12_resnet = []\n",
    "result_12_dp = []\n",
    "\n",
    "# 仿真\n",
    "result_12_resnet, result_12_dp = simulation_lips(embed_dim, num_heads,forward_factor,\n",
    "                                                 K,B,factor,episilon_factor,\n",
    "                                                 layers = [1,2,4,8])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12.792065,28.40365,58.267673,112.27357\n",
      "2.4456275,3.4856105,5.1531806,7.6673017\n"
     ]
    }
   ],
   "source": [
    "# 打印列表结果\n",
    "print(*np.array(result_12_resnet), sep=',')\n",
    "print(*np.array(result_12_dp), sep=',')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "d2l",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
